use std::{env, process::ExitCode};

use image::{ImageReader, RgbImage, Rgb32FImage, Rgb, error::ImageError, imageops::{self, FilterType}};

use rayon::iter::{ParallelIterator, IntoParallelIterator};

type Float = f32;

// 一维 DCT 表
pub struct DctTable {
    table: Vec<Vec<Float>>,
}

impl DctTable {
    pub fn new(length: usize) -> Self {
        let mut table = (0..length).into_par_iter().map(|i|{
            (i, Self::gen_dct_row(i as u32, length))
        }).collect::<Vec<(usize, Vec<Float>)>>();
        table.sort_by_key(|k| k.0);
        let table = table.into_iter().map(|(_, row)|{row}).collect::<Vec<Vec<Float>>>();
        Self {
            table,
        }
    }

    fn gen_dct_row(freq: u32, length: usize) -> Vec<Float> {
        const PI: f64 = std::f64::consts::PI;
        let mut row = Vec::<Float>::with_capacity(length);
        for i in 0..length {
            let p = (i as Float + 0.5) / length as Float;
            row.push((freq as Float * p * PI as Float).cos());
        }
        row
    }

    // 采样，从特定频率找特定位置的样本
    pub fn get(&self, freq: u32, sample: usize) -> Float {
        self.table[freq as usize][sample]
    }
}

// 二维 DCT 表
pub struct Dct2D {
    dct_row: DctTable,
    dct_col: DctTable,
}

impl Dct2D {
    pub fn new(width: usize, height: usize) -> Self {
        Self {
            dct_row: DctTable::new(width),
            dct_col: DctTable::new(height),
        }
    }

    // 采样。横竖的采样是相乘的。
    pub fn get(&self, freq_row: u32, sample_row: usize, freq_col: u32, sample_col: usize) -> Float {
        self.dct_row.get(freq_row, sample_row) * self.dct_col.get(freq_col, sample_col)
    }
}

// 用图片数据和二维 DCT 表做卷积，得到 (R, G, B)
fn dot_dct_kernel(img: &RgbImage, dct2d: &Dct2D, freq_h: u32, freq_v: u32) -> (Float, Float, Float) {
    let (w, h) = img.dimensions();
    let num_pixels = w * h;

    // 卷积矩阵的每一个元素
    let components = img.par_enumerate_pixels().map(|(x, y, p)| -> (Float, Float, Float) {
        let p = p.0;
        let (r, g, b) = (p[0], p[1], p[2]);

        // 原图处理到 [-1, 1] 范围
        let r = (r as Float / 255.0) * 2.0 - 1.0;
        let g = (g as Float / 255.0) * 2.0 - 1.0;
        let b = (b as Float / 255.0) * 2.0 - 1.0;

        let k = dct2d.get(freq_h, x as usize, freq_v, y as usize);

        // 卷积第一步：每个元素各自相乘
        (r * k, g * k, b * k)
    }).collect::<Vec<(Float, Float, Float)>>();

    // 卷积第二步：元素求和
    let (mut sr, mut sg, mut sb) = (0.0, 0.0, 0.0);
    for ttt in components.iter() {
        let (kr, kg, kb) = *ttt;
        sr += kr;
        sg += kg;
        sb += kb;
    }

    // 标准化处理
    sr /= num_pixels as Float;
    sg /= num_pixels as Float;
    sb /= num_pixels as Float;

    (sr, sg, sb)
}

// 跑 DCT 算法，把一张图变成一个 DCT 图。
fn run_dct(dct2d: &Dct2D, img: &RgbImage) -> Rgb32FImage {
    let (w, h) = img.dimensions();
    let mut dctimg = Rgb32FImage::new(w, h);
    dctimg.par_enumerate_pixels_mut().for_each(|(x, y, p)|{
        // 以 (x, y) 为频率值的二维DCT矩阵和原图做点乘
        let (sr, sg, sb) = dot_dct_kernel(&img, &dct2d, x, y);
        *p = Rgb([sr, sg, sb]);
    });
    dctimg
}

fn make_kernel_with_weights(w: u32, h: u32, dct2d: &Dct2D, freq_h: u32, freq_v: u32, dct_point: (Float, Float, Float)) -> Rgb32FImage {
    let mut image = Rgb32FImage::new(w, h);
    image.par_enumerate_pixels_mut().for_each(|(x, y, p)|{
        let k = dct2d.get(freq_h, x as usize, freq_v, y as usize);
        let (pr, pg, pb) = dct_point;
        let r = k * pr;
        let g = k * pg;
        let b = k * pb;
        *p = Rgb([r, g, b])
    });
    image
}

fn add_idct_images(add_to: &mut Rgb32FImage, source: &Rgb32FImage) {
    assert_eq!(add_to.dimensions(), source.dimensions());

    add_to.par_enumerate_pixels_mut().for_each(|(x, y, p)|{
        let ps = source.get_pixel(x, y).0;
        let pd = p.0;
        let (sr, sg, sb) = (ps[0], ps[1], ps[2]);
        let (dr, dg, db) = (pd[0], pd[1], pd[2]);
        *p = Rgb([sr + dr, sg + dg, sb + db]);
    });
}

fn run_idct(dct2d: &Dct2D, dctimg: &Rgb32FImage) -> RgbImage {
    let (w, h) = dctimg.dimensions();
    let mut idctimg = dctimg.clone();

    dctimg.enumerate_pixels().for_each(|(x, y, p)|{
        let p = p.0;
        let point = (p[0], p[1], p[2]);
        let layer = make_kernel_with_weights(w, h, dct2d, x, y, point);
        add_idct_images(&mut idctimg, &layer);
    });

    let mut saveimg = RgbImage::new(w, h);
    saveimg.par_enumerate_pixels_mut().for_each(|(x, y, p)|{
        let idctp = idctimg.get_pixel(x, y).0;
        let (r, g, b) = (idctp[0], idctp[1], idctp[2]);
        let r = r * 0.5 + 0.5;
        let g = g * 0.5 + 0.5;
        let b = b * 0.5 + 0.5;
        let r = r.clamp(0.0, 1.0);
        let g = g.clamp(0.0, 1.0);
        let b = b.clamp(0.0, 1.0);
        let r = (r * 255.0) as u8;
        let g = (g * 255.0) as u8;
        let b = (b * 255.0) as u8;
        *p = Rgb([r, g, b]);
    });
    saveimg
}

fn save_dct_for_visualize(dctimg: &Rgb32FImage, save_to: &str) -> Result<(), ImageError> {
    let (w, h) = dctimg.dimensions();
    let mut saveimg = RgbImage::new(w, h);

    saveimg.par_enumerate_pixels_mut().for_each(|(x, y, p)|{
        let dctp = dctimg.get_pixel(x, y).0;
        let (r, g, b) = (dctp[0], dctp[1], dctp[2]);
        let r = (r * 0.5 + 0.5).clamp(0.0, 1.0);
        let g = (g * 0.5 + 0.5).clamp(0.0, 1.0);
        let b = (b * 0.5 + 0.5).clamp(0.0, 1.0);
        let r = (r * 255.0) as u8;
        let g = (g * 255.0) as u8;
        let b = (b * 255.0) as u8;
        *p =Rgb([r, g, b]);
    });
    saveimg.save(save_to)?;
    Ok(())
}

// 做测试，把一张图变成一个 DCT 图，再利用 DCT 图把它变回来。
fn do_test(src_path: &str, hdr_path: &str, idct_path: &str, scaling: f64) -> Result<(), ImageError> {
    let srcimg = ImageReader::open(src_path)?.decode()?.to_rgb8();
    let (w, h) = srcimg.dimensions();
    let dct2d = Dct2D::new(w as usize, h as usize);

    // 先跑 DCT
    println!("Running DCT {src_path} -> {hdr_path}");
    let dctimg = run_dct(&dct2d, &srcimg);
    dctimg.save(hdr_path)?;

    // 显示可视化的 DCT 图
    save_dct_for_visualize(&dctimg, &format!("{}.png", hdr_path.rsplitn(2, ".").collect::<Vec<&str>>()[1]))?;

    // 再跑 IDCT
    println!("Running IDCT {hdr_path} -> {idct_path}");

    // 先进行缩放
    let new_w = ((w as f64) * scaling) as u32;
    let new_h = ((h as f64) * scaling) as u32;
    let new_dct2d = Dct2D::new(new_w as usize, new_h as usize);
    let new_dctimg = imageops::resize(&dctimg, new_w, new_h, FilterType::Triangle);

    let saveimg = run_idct(&new_dct2d, &new_dctimg);
    saveimg.save(idct_path)?;

    Ok(())
}

fn usage() {
    println!("Usage: dstscaling <input.png> <out.hdr> <out.png> <scaling>");
    println!("Can also use JPEG or other formats.");
}

fn main() -> ExitCode {
    let args: Vec<String> = env::args().collect();
    if args.len() < 5 {usage(); return ExitCode::from(1);}
    let input_file = args[1].clone();
    let hdr_file = args[2].clone();
    let idct_file = args[3].clone();
    let scaling = match args[4].parse::<f64>() {
        Ok(scaling) => scaling,
        Err(err) => {
            eprintln!("{:?}, bad param, got `{}`", err, args[3]);
            usage();
            return ExitCode::from(1);
        },
    };

    match do_test(&input_file, &hdr_file, &idct_file, scaling) {
        Ok(_) => ExitCode::from(0),
        Err(err) => {
            eprintln!("{:?}", err);
            ExitCode::from(2)
        },
    }
}
